import glob
import os
import os.path as osp
import cv2
import numpy as np
import pdb
import sys
import open3d as o3d

import cv2
import sys
from get_cylinder_xyz import *
camera_type = ['frontleft','frontmid','frontright','rearleft','rearright', 'topleft','topright',"stereoleft","stereoright"]
camera_type_w_name_dict_ = {"front_left":'frontleft',
                            "front_mid":'frontmid',
                            "front_right":'frontright',
                            "rear_left":'rearleft',
                            "rear_right":'rearright',
                            "stereo_left":'stereoleft',
                            "stereo_right":'stereoright',
                            "top_left":'topleft',
                            "top_right":'topright'}
def mkdir_or_exist(dir_name, mode=0o777):
    if dir_name == '':
        return
    dir_name = osp.expanduser(dir_name)
    os.makedirs(dir_name, mode=mode, exist_ok=True)

def showImg(img,wati_time=0,win_name="r"):
    print(img.shape)
    cv2.namedWindow(win_name,cv2.WINDOW_NORMAL)
    cv2.imshow(win_name,img)
    key = cv2.waitKey(wati_time) & 0xFF
    if key == ord('q'):
        sys.exit()
    return key


def get_KE(K_p,E_p,side_add_w=0,use_1400 = True):
    fs  = cv2.FileStorage(K_p, cv2.FileStorage_READ)
    K = fs.getNode('K').mat()
    D = fs.getNode('D').mat()
    
    if use_1400:
        K = K/2
        K[0][2]+=side_add_w/2
        K[2][2]=1
    if osp.exists(E_p)==False:
        print("file not exist: \n",E_p)
        exit()
    E = np.loadtxt(E_p)
    return K,D,E

def get_img_pts(img):
    cv2.namedWindow('image',cv2.WINDOW_NORMAL)

    selected_point = []

    def get_coordinate(event, x, y, flags, param):
        # global selected_point
        if event == cv2.EVENT_LBUTTONDBLCLK:
            cv2.circle(img,(x,y),1,(0,255,0),-1)
            selected_point.append([x,y])

    cv2.setMouseCallback('image', get_coordinate)

    while True:
        cv2.imshow('image', img)
        if cv2.waitKey(20) & 0xFF == 27:
            break

    cv2.destroyAllWindows()
    return np.array(selected_point)


if __name__=="__main__":

    img_p=sys.argv[1]
    lidar_p=sys.argv[2]
    qtruck_name = sys.argv[3] # qtruck36
    camera_name=sys.argv[4]  #front_right
    # use_undis=sys.argv[5] # false
    
    xyz_min=[6, -2, 0.2]
    xyz_max=[25,2, 1.5]
    
    use_undis="false"

    # stereo=sys.argv[6]
    stereo="false"
    if stereo == "true":
        side_add_w=0
        use_1400 = False
    else:
        side_add_w=440
        use_1400 = True
        
        side_add_w=0
        use_1400 = False
        
    K_root = "/home/westwell/welldriver/wellpilot_config/calib/cameras_in_calib_file"
    E_root = "/home/westwell/welldriver/wellpilot_config/calib/cameras_ex_calib_file"

    K_p=osp.join(K_root,qtruck_name,camera_name+".yaml")
    E_p=osp.join(E_root,qtruck_name,camera_name+".txt")
    flags = [osp.exists(img_p),osp.exists(lidar_p),osp.exists(K_p),osp.exists(E_p)]

    if osp.exists(E_p)==False:
        print("file not exist:, create I \n",E_p)
        mkdir_or_exist(osp.dirname(E_p))
        np.savetxt(E_p,np.eye(4))
        
        
    for p,flag in zip([img_p,lidar_p,K_p,E_p],flags):
        print(flag,p)
        

    K,D,E=get_KE(K_p,E_p,side_add_w,use_1400)
    print("K \n",K)

    print("use save txt or not :")
    # use_txt=input()
    use_txt="0"
    if use_txt == "1":
        lidar_pts=np.loadtxt("lidar_pts.txt",dtype=np.float32)
        img_pts=np.loadtxt("img_pts.txt",dtype=np.float32)
        print(lidar_pts)
        print(img_pts)

    else:


        # print("please input pts num : ")
        # pts_num = int(input()) # 4 
        auto_det_p_lidar = False
        debug= True
        pts_num = 4
        if auto_det_p_lidar:
            lidar_pts = get_lidar_pts_auto(lidar_p,pts_num,xyz_min,xyz_max,debug)
        else:
            lidar_pts = get_lidar_pts_manual(lidar_p,pts_num)
            # lidar_pts = get_lidar_pts_manual_not_fit(lidar_p,pts_num)
            
        print(lidar_pts)
        np.savetxt("lidar_pts.txt",lidar_pts)

        img = cv2.imread(img_p)
        if use_undis=="true":
            img = cv2.undistort(img,K,D)
        img_pts = get_img_pts(img)
        img_pts = img_pts.astype(np.float32)
        print(img_pts)
        np.savetxt("img_pts.txt",img_pts)

    retval, rvec, tvec = cv2.solvePnP(lidar_pts, img_pts, K, None)
    R, _ = cv2.Rodrigues(rvec)
    P = np.concatenate([R,tvec],axis=1)
    P = np.vstack([P,[0,0,0,1]])
    for p in P:
        print(p[0],p[1],p[2],p[3])

    np.savetxt(E_p,P,fmt='%.8f')
    # np.savetxt(E_p,P)
    print(E_p)